feat(executorch): expose composable Edge export API - #4440
Conversation
cd55b95 to
4dc2d25
Compare
4dc2d25 to
cb3633f
Compare
| graph_module: torch.fx.GraphModule, | ||
| payload_memo: dict[int, Any], | ||
| ) -> torch.fx.GraphModule: | ||
| staged = copy.deepcopy(graph_module, payload_memo) |
There was a problem hiding this comment.
Does this doubles the GPU/Host memory usage? Is the copy absolutely necessary?
There was a problem hiding this comment.
Or is there a way to not copy the weight, only the graph structure
There was a problem hiding this comment.
Good question to push on, and no, weights are not copied. Only graph structure is,
which is what you suggest in your next comment.
_payload_sharing_memo (just above) pre-seeds copy.deepcopy's memo with every tensor
and engine payload reachable from the program, so deepcopy returns those objects
as-is instead of duplicating them. state_dict and constants are then shallow-copied
dicts, so the keys are independent but the values are the same objects.
I measured it to be sure, on a model with 33.6 MB of weights:
weights in program: 33.6 MB
live CPU tensor bytes before staging: 33.6 MB
live CPU tensor bytes after staging: 33.6 MB
growth: 0.0 MB
state_dict entries: 4, shared by identity: 4
a.weight same object: True same storage: True
a.bias same object: True same storage: True
b.weight same object: True same storage: True
b.bias same object: True same storage: True
So the copy is graph nodes, signatures, and metadata containers, all small. TensorRT
engine blobs are shared the same way, which matters more since those can be multiple
gigabytes.
On whether the copy is necessary at all: replace_execute_engine rewrites the graph in
place (registers a buffer, erases nodes, edits the graph signature). Without staging, a
failure part-way through would leave the caller's ExportedProgram half-rewritten. The
GraphModule path happens to be safe already because dynamo/_exporter.py deep-copies,
but the pre-exported ExportedProgram and multi-method paths are handed to us directly,
so staging is what keeps those inputs intact.
The cost this leaves is documented as a constraint: transform passes must not mutate the
shared payloads in place, since those are the caller's tensors.
There was a problem hiding this comment.
Yes, that is exactly what it does. Only graph structure is copied; tensors and engine
payloads are shared by identity through the pre-seeded deepcopy memo. Measurements in
the reply above: 33.6 MB of weights in, 0.0 MB growth, all four state_dict tensors
identical objects sharing the same storage.
One case needed a carve-out, which is worth mentioning since it was a real bug. Metadata
describing symbolic shapes (SymInt, fake tensors) cannot be deep-copied at all: it
reaches back into the live ShapeEnv, which holds fake tensors, and copying raised
RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor)
on every dynamic-shape model. That metadata is now shared rather than copied, which also
keeps the staged graph attached to the shape symbols it is guarded on. Covered by
test_stage_exported_program_supports_dynamic_shapes.
| graph_module.register_buffer(buffer_name, engine_tensor, persistent=True) | ||
| exported_program.state_dict[buffer_name] = engine_tensor | ||
| str_args = [ | ||
| str(value) if value is not None else "" |
There was a problem hiding this comment.
engine_info[ENGINE_IDX] is raw bytes (often hundreds of MB / multi-GB). str(b"...") builds a temporary Python string of roughly 4 * len(bytes) characters (b'\xab\xcd...' style), then immediately throws that slot away and uses engine_attr_node instead.
There was a problem hiding this comment.
Good catch, this was a real waste and it is fixed now.
You are right that the slot is thrown away: no_op_args splices in engine_attr_node at
ENGINE_IDX, so the stringified engine was built and then never used. Measured on a
64 MB stand-in payload:
engine bytes : 67.1 MB
str(bytes) length : 268.4 M chars
str object bytes : 268.4 MB
blowup factor : 4.0 x
So a 2 GB engine burned about 8 GB of host memory building a string that was immediately
discarded.
Now that index is skipped:
str_args = [
("" if index == ENGINE_IDX else str(value) if value is not None else "")
for index, value in enumerate(engine_info[:SERIALIZATION_LEN])
]Re-verified on a GPU afterward: static-shape and dynamic-shape export() still produce
loadable .pte files of the same size as before, and save(output_format="executorch")
is unchanged.
| shared mutable state between them. | ||
| """ | ||
| from torch_tensorrt._features import ENABLED_FEATURES | ||
|
|
There was a problem hiding this comment.
Do we need a check that the platform is linux?
There was a problem hiding this comment.
Yes, we do, and it was missing. Added.
save(output_format="executorch") already rejects non-Linux platforms in _compile.py,
so the advanced entry point was skipping a check the one-step path enforces. export()
now raises the same error with the same wording:
if platform.system() != "Linux":
raise ValueError(
f"The executorch format is only supported on Linux, {platform.system()} "
"is not a supported platform for this format"
)Placed before the runtime-feature check so the platform message comes first, and covered
by test_export_rejects_non_linux_platform.
## What this adds
`torch_tensorrt.save(..., output_format="executorch")` turns a compiled model
into an ExecuTorch `.pte` file in one step. That is the easiest path and it stays
the recommendation.
Some workflows need to stop earlier. You may want to look at the delegated Edge
graph, run your own Edge transforms, send operations TensorRT cannot handle to
another backend, keep separate methods such as `prefill` and `decode`, add
constant methods, or pick the final ExecuTorch configuration yourself.
This adds `torch_tensorrt.executorch.export()`. It returns ExecuTorch's standard
`EdgeProgramManager`, which is the supported place to inspect and customize an
Edge program before you call `to_executorch()`.
```python
import torch_tensorrt.executorch
edge = torch_tensorrt.executorch.export(trt_module)
# Inspect or transform the Edge program here.
print(edge.exported_program().graph)
program = edge.to_executorch()
with open("model.pte", "wb") as output:
program.write_to_file(output)
# Advanced callers must also persist external tensor data when present.
program.write_tensor_data_to_file(".")
```
`save()` now calls the same implementation, so the simple and advanced paths
share one lowering path instead of two.
## Which one should I use?
Use `save()` when you just want a `.pte`. Use
`torch_tensorrt.executorch.export()` when you need control before the `.pte` is
created.
## Example: TensorRT with a fallback backend
TensorRT always gets the first chance to claim its prebuilt engine nodes. Any
partitioners you pass run afterward and can claim operations TensorRT does not
support.
```python
edge = torch_tensorrt.executorch.export(
trt_module,
partitioners=[cuda_partitioner],
)
program = edge.to_executorch()
```
## Example: multiple methods
```python
edge = torch_tensorrt.executorch.export(
{
"prefill": prefill_program,
"decode": decode_program,
},
compile_specs={
"prefill": prefill_specs,
"decode": decode_specs,
},
)
```
Method names are preserved, and each method can have its own partitioner and
compile-spec pipeline. This does not make mutable state shared between methods.
## Accepted inputs
- a TensorRT-compiled `torch.fx.GraphModule`
- an engine-bearing `torch.export.ExportedProgram`
- a mapping from method names to independent `ExportedProgram` objects
A plain `torch.nn.Module` must be compiled with Torch-TensorRT first.
## Safety and memory behavior
Rewriting engine calls changes the exported graph. To avoid damaging the program
you passed in, export first stages its own copy of the graph structure,
signatures, state containers, constants, and metadata. If rewriting or lowering
fails, your original program is left alone.
Tensor and TensorRT engine payloads are treated as read-only and shared with
that staged copy, so a multi-gigabyte engine is not duplicated. Custom transform
passes must not modify these shared payload objects.
Metadata that describes symbolic shapes is a special case. It points back to the
live shape environment that the exported program is guarded on, so it is shared
rather than copied. Copying it would both fail (the shape environment holds fake
tensors that cannot be duplicated) and detach the copy from the symbols the graph
depends on. This is what keeps dynamic-shape models working.
When a lifted TensorRT engine is replaced, the old placeholder, its
graph-signature entry, and its constant are removed. Several calls that share one
engine reuse a single materialized `uint8` payload buffer.
## Compatibility
- Existing one-step `save(..., output_format="executorch")` behavior is preserved,
including its rejection of non-list `partitioners` and `compile_specs`.
- `"executorch"` is now included in the `output_format` type hint. It was already
accepted at run time, so this only fixes the annotation.
- `.pte` serialization and external `.ptd` persistence are preserved.
- TensorRT remains the first partitioner.
- Shared dynamic dimensions declared with `Input(shared_dims=...)` remain shared.
- Zero-engine methods are allowed. A later partitioner may claim them, or portable
operators may remain undelegated.
## Testing
Tested on a CUDA GPU with TensorRT and ExecuTorch installed, building
Torch-TensorRT from source so the C++ runtime was active.
Verified by compiling a real model with TensorRT and then exporting it:
- static-shape `export()` produces a loadable `.pte`
- dynamic-shape `export()` produces a loadable `.pte`, matching the output of the
previous one-step save path
- the source `GraphModule` still holds its original engine nodes afterward
- engine payloads are shared, not re-serialized
- `save(..., output_format="executorch")` still writes a `.pte`
- `save()` still raises `TypeError` for a non-list `partitioners` or
`compile_specs`
Also ran the ExecuTorch test directory (76 passed). Two GPU tests in
`test_cuda_partitioner_composition.py` fail, and they fail the same way with this
change reverted, so they are not caused by it.
Static checks: formatting, import ordering, lint, and type checking on the changed
files.
cb3633f to
93efd54
Compare
Why this PR is needed
torch_tensorrt.save(..., output_format="executorch")is the easiest way to create an ExecuTorch.ptefile. It completes the whole export pipeline immediately:That remains the recommended API for common one-step exports.
Some workflows need to stop before final ExecuTorch memory planning and serialization. For example, users may want to:
prefillanddecode;This PR adds
torch_tensorrt.executorch.export(). It returns ExecuTorch's standardEdgeProgramManager, which is the supported boundary for inspecting and customizing an Edge program before callingto_executorch().Which API should I use?
Use
torch_tensorrt.save()when you want a.ptein one step.Use
torch_tensorrt.executorch.export()when you need control before the final.pteis created:save()now delegates to the same implementation, so the simple and advanced paths do not maintain separate lowering logic.Example: TensorRT with a fallback backend
TensorRT always receives the first chance to claim its prebuilt engine nodes. Additional partitioners run afterward and can claim operations that TensorRT does not support.
Assume
cuda_partitioneris an ExecuTorch partitioner configured for your target:This can produce one ExecuTorch program containing both TensorRT and CUDA delegates.
Example: multiple methods
Assume
prefill_programanddecode_programare independent, engine-bearingExportedProgramobjects, and the two spec lists are configured for their methods:The mapping preserves method names and allows per-method partitioner and compile-spec pipelines. It does not imply that mutable state is shared between methods.
Example: transforms, metadata, and ETRecord
Assume the transform passes and Edge configuration below are provided by the caller:
Accepted inputs
The new API accepts:
torch.fx.GraphModule;torch.export.ExportedProgram;ExportedProgramobjects.A plain
torch.nn.Modulemust be compiled with Torch-TensorRT first.Safety and memory behavior
Engine rewriting changes the exported graph. The new API stages an independent copy of graph structure, signatures, state containers, constants, and metadata before rewriting anything. Rewrite and lowering failures therefore do not consume or structurally modify the caller's original program.
Tensor and TensorRT engine payload storage is treated as immutable and shared with the staged structure. This avoids deep-copying potentially multi-gigabyte engines. Custom transform passes must not mutate shared payload objects.
Metadata describing symbolic shapes is also shared rather than copied. It points back into the live shape environment the exported program is guarded on, and that environment holds fake tensors which cannot be copied. Sharing it is what keeps dynamic-shape models working; copying it would also detach the copy from the shape symbols the graph depends on.
When a lifted TensorRT engine is replaced, the old placeholder, graph-signature entry, and constant are removed. Multiple calls that share one engine reuse one materialized
uint8payload buffer.Compatibility
save(..., output_format="executorch")behavior is preserved..pteserialization and external.ptdpersistence are preserved.Input(shared_dims=...)remain shared.Validation
Built from source with the Torch-TensorRT C++ runtime active, against TensorRT and
ExecuTorch, and exercised on both x86_64 and aarch64 CUDA GPUs.
Checked end to end with a real compiled TensorRT engine:
export()writes a loadable.pte;export()writes a loadable.pteof the same size the previousone-step save path produced;
GraphModulestill holds its original engine nodes afterexport();save(..., output_format="executorch")still writes a.pte;save()still raisesTypeErrorfor a non-listpartitionersorcompile_specs.Unit and integration tests in the ExecuTorch test directory: 76 passed.
Two GPU tests in
test_cuda_partitioner_composition.py,test_erfinv_routes_to_cuda_backendandtest_weighted_partition_persists_external_data, fail withSpecViolationError: Node.meta lowered_module_0 is missing val field. They fail thesame way on main without this change, so they are pre-existing and not caused by
this work.
Note that these tests only run in the L2 tier, which requires the
ci: run-l2labelon a pull request.
Static checks on the changed files: import ordering, formatting, lint, and type
checking.
Related work
PR #4433 forwards additional Edge-lowering options through one-step
save(). After this PR lands, #4433 can keepsave()as a thin wrapper overtorch_tensorrt.executorch.export()instead of extending the previous private lowering path.